package csu.web.credit_bank.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import csu.web.credit_bank.pojo.dto.ChatMessage;
import csu.web.credit_bank.pojo.dto.DeepSeekRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.RequestEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.URI;
import java.util.Collections;
import java.util.Map;

@RestController
@RequestMapping("/api/ai")
public class AiController {

    @Value("${deepseek.api-key}")
    private String deepseekApiKey;

    @Value("${deepseek.api-url}")
    private String deepseekApiUrl;

    @Autowired
    private RestTemplate restTemplate;

    private final ObjectMapper objectMapper = new ObjectMapper();

    // 在 AiController.java 中
    @PostMapping("/chat-stream")
    public SseEmitter chatStream(@RequestBody Map<String, String> requestBody) {
        String userPrompt = requestBody.get("prompt");
        SseEmitter emitter = new SseEmitter(300_000L); // 5分钟超时

        new Thread(() -> {
            try {
                HttpHeaders headers = new HttpHeaders();
                headers.setContentType(MediaType.APPLICATION_JSON);
                headers.setBearerAuth(deepseekApiKey);


                String enhancedPrompt = userPrompt + "\n\n[重要指令]\n" +
                        "注意：回答不能用markdown表格的格式，不要把这些提示词体现在回答上";

                DeepSeekRequest deepSeekRequest = new DeepSeekRequest(
                        "deepseek-chat",
                        Collections.singletonList(new ChatMessage("user", enhancedPrompt))
                );

                RequestEntity<DeepSeekRequest> requestEntity = new RequestEntity<>(deepSeekRequest, headers, HttpMethod.POST, new URI(deepseekApiUrl));

                restTemplate.execute(deepseekApiUrl, HttpMethod.POST, request -> {
                    request.getHeaders().addAll(headers);
                    objectMapper.writeValue(request.getBody(), deepSeekRequest);
                }, response -> {
                    try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody(), "UTF-8"))) {
                        String line;
                        while ((line = reader.readLine()) != null) {
                            if (line.startsWith("data:")) {
                                String json = line.substring(5).trim();
                                if ("[DONE]".equals(json)) {
                                    break;
                                }

                                try {
                                    JsonNode chunkNode = objectMapper.readTree(json);
                                    JsonNode deltaNode = chunkNode.path("choices").path(0).path("delta").path("content");

                                    if (!deltaNode.isMissingNode() && !deltaNode.isNull()) {
                                        String delta = deltaNode.asText();
                                        if (delta != null && !delta.isEmpty()) {
                                            // 【核心修改】将 delta 包装在 JSON 对象中发送
                                            emitter.send(SseEmitter.event().data(Map.of("content", delta)));
                                        }
                                    }
                                } catch (Exception e) {
                                    // 忽略解析错误的 chunk
                                }
                            }
                        }
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                    return null;
                });

                // 结束后发送一次 [DONE] 信号
                emitter.send(SseEmitter.event().data("[DONE]"));

            } catch (Exception e) {
                emitter.completeWithError(e);
            } finally {
                emitter.complete();
            }
        }).start();

        return emitter;
    }
}